import numpy as np

class LVQ:
    # 计算两个向量的欧氏距离
    def euclidean_dist(self, x1, x2):
        return np.linalg.norm(x1 - x2)  # 求范数函数

    # 选择与样本最近的原型向量的编号
    def get_min_dist_proto_vect(self, proto_vects, x):
        min_dist = np.inf
        min_index = -1
        for i in range(len(proto_vects)):
            dist = self.euclidean_dist(proto_vects[i], x)
            if dist < min_dist:
                min_dist = dist
                min_index = i
        return min_index

    # LVQ训练
    def fit(self, data_x, data_y, proto_vects, proto_labels, learnrate, max_iters):
        """
        data_x: 训练集特征数据
        data_y: 训练样本的类别标记
        proto_vects: 初始化的原型向量
        proto_labels: 原型向量对应的类别标记
        learnrate:更新速率
        max_iters: 最大迭代次数
        """
        iters = 0
        while iters < max_iters:
            # 随机选择一个训练样本sample_index
            sample_index = np.random.choice(len(data_x), 1, replace=False)
            min_dist_index = self.get_min_dist_proto_vect(proto_vects, data_x[sample_index])
            # 根据样本的类别来更新原型向量
            if proto_labels[min_dist_index] == data_y[sample_index]:
                proto_vects[min_dist_index] = proto_vects[min_dist_index] + \
                                              learnrate * (data_x[sample_index] - proto_vects[min_dist_index])
            else:
                proto_vects[min_dist_index] = proto_vects[min_dist_index] - \
                                              learnrate * (data_x[sample_index] - proto_vects[min_dist_index])
            iters += 1


if __name__ == '__main__':
    from sklearn.datasets._samples_generator import make_blobs
    import matplotlib.pyplot as plt
    from random import randint

    fig = plt.figure(1)
    plt.subplot(2, 1, 1)
    center = [[1, 1], [-1, -1], [1, -1]]
    cluster_std = 0.35
    X1, Y1 = make_blobs(n_samples=100000, centers=center, n_features=2, cluster_std=cluster_std, random_state=1)
    plt.scatter(X1[:, 0], X1[:, 1], marker='o', c=Y1)

    def random_select_k_samples(sample_num, k):
        sample_index = []
        while len(sample_index) < k:
            index = randint(0, sample_num - 1)
            if index in sample_index:
                continue
            else:
                sample_index.append(index)
        return sample_index

    plt.subplot(2, 1, 2)

    sample_index = random_select_k_samples(len(X1), 3)
    proto_vects = X1[sample_index]
    proto_labels = Y1[sample_index]
    lvq = LVQ()
    lvq.fit(X1, Y1, proto_vects, proto_labels, 0.04, 10000)
    mus = proto_vects
    plt.scatter(X1[:, 0], X1[:, 1], marker='o', c=Y1)
    plt.scatter(mus[:, 0], mus[:, 1], marker='^', c='r')
    plt.savefig("figure.jpg")